import json
import re
from typing import Literal, Optional

from pydantic import BaseModel
from sse_starlette import EventSourceResponse

from utils.ai_chess.ai_chess import ai_chess


class Message(BaseModel):
    role: Literal['user', 'ai']
    text: Optional[str] = None
    board: Optional[str] = None
    id: Optional[str] = None
    f: Optional[str] = None
    t: Optional[str] = None


class InputParams(BaseModel):
    historyMessage: list[Message]


async def ai_move(params: InputParams):
    async def chat_event():
        chat_history = []
        for history in params.historyMessage:
            message = {
                'content': {
                    'id': history.id,
                    'from': history.f,
                    'to': history.t
                }
            }
            if history.role == 'user':
                message['role'] = 'user'
                message['content']['text'] = history.text
                message['content']['board'] = history.board
            elif history.role == 'ai':
                message['role'] = 'assistant'
                if history.text:
                    message = {
                        'role': 'assistant',
                        'content': {
                            'text': history.text
                        }
                    }
            message['content'] = json.dumps(message['content'], ensure_ascii=False)
            chat_history.append(message)

        mode = 'json'
        json_string = ""
        start = False
        answer = ai_chess(chat_history)
        for a in answer:
            if a == "<think>":
                mode = 'think'
                continue
            elif a == '</think>':
                mode = 'json'
                continue

            if mode == 'think':
                if not start:
                    if a == "\n":
                        start = True
                        continue
                yield json.dumps({'type': 'think', 'data': a}, ensure_ascii=False)
            else:
                json_string += a

        json_match = re.findall(r'```json?(.*?)```', json_string, re.DOTALL)
        if json_match:
            json_string = json_match[0]
            data = json.loads(json_string)
            fx = int(data['from'][1]) - 1
            fy = ord(data['from'][0]) - 65
            tx = int(data['to'][1]) - 1
            ty = ord(data['to'][0]) - 65
            data = {
                'type': 'message',
                'data': {
                    'role': 'ai',
                    'id': data['id'],
                    'from': {
                        'x': fx,
                        'y': fy,
                        'string': data['from']
                    },
                    'to': {
                        'x': tx,
                        'y': ty,
                        'string': data['to']
                    }
                }
            }

            yield json.dumps(data, ensure_ascii=False)

    return EventSourceResponse(chat_event())
